import numpy as np
from .base import BinadPlainBuilder


# @njit(parallel=True)
# def adjugate_matrix(matrix):
#     """
#     计算矩阵的伴随矩阵
#     """
#     n = matrix.shape[-1]
#     m = matrix.shape[:-3]
#     adjugate = np.zeros_like(matrix)

#     for i in prange(n):
#         for j in prange(n):
#             # Create the minor matrix by removing the i-th row and j-th column
#             minor = np.zeros((*m, n - 1, n - 1))
#             for k in range(n):
#                 if k == i:
#                     continue
#                 for l in range(n):
#                     if l == j:
#                         continue
#                     minor_k = k if k < i else k - 1
#                     minor_l = l if l < j else l - 1
#                     minor[..., minor_k, minor_l] = matrix[..., k, l]
#             # Calculate the cofactor
#             cofactor = (-1) ** (i + j) * np.linalg.det(minor)
#             adjugate[..., j, i] = cofactor  # Note the transpose here


#     return adjugate
def adjugate_matrix(matrix):
    """
    计算矩阵的伴随矩阵，后面再考虑如何并行加速之类的，或者压根就影响不大
    """
    assert matrix.shape[-1] == matrix.shape[-2]
    n = matrix.shape[-1]
    adjugate = np.zeros_like(matrix)

    for i in range(n):
        for j in range(n):
            # Create the minor matrix by removing the i-th row and j-th column
            minor = np.delete(np.delete(matrix, i, axis=-2), j, axis=-1)
            # Calculate the cofactor
            cofactor = (-1) ** (i + j) * np.linalg.det(minor)
            adjugate[..., j, i] = cofactor  # Note the transpose here

    return adjugate


class ChristianBinadBuilder(BinadPlainBuilder):
    """
    按论文[1] Christian J. A. , Derksen H. , Watkins R. .Lunar crater identification in digital Images[J/OL].J. Astronaut. Sci.,2021,68(4):1056-1144
    中的不共面不变量构建四条公切线的两个交点计算不变量
    """

    def invariant(self, C1, C2):
        """
        本函数需要并行化计算
        """
        assert C1.shape[1] == 3 and C2.shape[1] == 3
        # 计算公切线的交点，由于numpy没有提供直接求解伴随矩阵的方法，而椭圆矩阵是满秩的，因此直接求逆就行了
        inv_C1 = np.linalg.inv(C1)
        inv_C2 = np.linalg.inv(C2)
        A = -inv_C1 @ C2
        # 作特征值分解
        eig_vals, _ = np.linalg.eig(A)
        index = np.isreal(eig_vals).all(axis=1)
        ## 排除复数特征值（TODO ？？ 研究一下，相当于排除相切的情况）
        eig_vals = eig_vals[index].real[..., None, None]
        inv_C1 = inv_C1.repeat(index.sum(), axis=0)[:, None]
        inv_C2 = inv_C2[index, None]
        # 构造退化二次型
        B = eig_vals * inv_C1 + inv_C2
        B_adj = -adjugate_matrix(B)
        ## 找出最大的主对角线元素对应的索引
        ## TODO 确认一下，是不是最大的对角线元素始终位于最后一项 N * 3
        max_diag_ind = np.abs(B_adj.diagonal(axis1=2, axis2=3)).argmax(axis=2)
        ## 对应的最大对角线元素 N*3
        diag = np.take_along_axis(
            B_adj.diagonal(axis1=2, axis2=3), max_diag_ind[..., None], axis=2
        ).squeeze()
        ## 只需要找出为负的对角线元素，且以最小的为准
        index = diag > 0
        ## 先不考虑负数问题，直接取绝对值
        z = (
            np.take_along_axis(B_adj, max_diag_ind[:, None, None], axis=3)[..., 0]
            / np.sqrt(np.abs(diag))[..., None]
        )
        ## 再排除对角线为负的情况:  z[index]
        ## 再排除两个解以上的情况
        ## 求解反对称矩阵
        Zx = np.array(
            [
                [np.zeros((index.sum())), -z[index, 2], z[index, 1]],
                [z[index, 2], np.zeros((index.sum())), -z[index, 0]],
                [-z[index, 1], z[index, 0], np.zeros((index.sum()))],
            ]
        ).transpose(2, 0, 1)
        # 求解两条交叉直线
        D = B[index] + Zx
        ## 直接取第一行和第一列作为g和h
        g = D[:, 0]
        h = D[:, :, 0]
        test = np.einsum("ni,nj->nij", g, h)
        p1 = z @ inv_C1
        p2 = z @ inv_C2
        # 计算不变量
        I1 = (
            np.einsum("in,nij,jn->n", p1, C1, p1)
            * np.einsum("in,nij,jn->n", p2, C1, p2)
            / np.einsum("in,nij,jn->n", p1, C1, p2) ** 2
        )
        I2 = (
            np.einsum("in,nij,jn->n", p2, C2, p2)
            * np.einsum("in,nij,jn->n", p1, C2, p1)
            / np.einsum("in,nij,jn->n", p1, C2, p2) ** 2
        )
        return I1, I2
